#!/usr/bin/env python
# coding: utf-8

import pandas as pd
import numpy as np
import os
import glob
from tqdm import tqdm
import warnings
from datetime import datetime, timedelta
import multiprocessing as mp
import sys
import time
import matplotlib.pyplot as plt
import seaborn as sns

# --- Matplotlib and Seaborn Settings (English Labels) ---
# plt.rcParams['font.sans-serif'] = ['Noto Sans CJK SC'] # Commented out - Use English
plt.rcParams['axes.unicode_minus'] = False
sns.set_theme(style="whitegrid")

# --- 配置参数 ---
# 数据路径 (需要与 task2/task3 一致)
# BASE_DATA_DIR = r"C:\Users\ASUS\Desktop\B题-全部数据" #<--- Windows 路径示例
BASE_DATA_DIR = "/share/home/u2201230214/GPR/B" #<--- Linux 路径
METADATA_TEST = os.path.join(BASE_DATA_DIR, "附件2", "Metadata2.csv")

# 任务二结果路径 (用于读取 MET 预测)
# TASK2_PREDICTIONS_DIR = "result_2" #<--- 相对路径示例
TASK2_PREDICTIONS_DIR = "/share/home/u2201230214/GPR/result_2" #<--- Linux 绝对路径

# 输出路径 (符合比赛要求)
# 先定义结果目录
RESULTS_DIR = "results_task4" #<--- 新增：任务四结果目录
# 然后使用它来构建其他路径
OUTPUT_SUMMARY_FILE = os.path.join(RESULTS_DIR, "result_4_summary.xlsx") # Changed filename slightly for clarity
FIGURES_SAVE_PATH = os.path.join(RESULTS_DIR, "figures") #<--- 新增：图表保存路径

# 久坐行为定义 (来自 TDB.md)
MET_STATIC_THRESHOLD = 1.6 # MET 低于此阈值为静态
MET_SLEEP_THRESHOLD = 1.0 # MET 低于此阈值视为睡眠 (用于排除)
SEDENTARY_DURATION_MINUTES = 30 # 连续静态超过此时长为久坐
SEDENTARY_DURATION_SECONDS = SEDENTARY_DURATION_MINUTES * 60

# 滑动窗口参数 (必须与 task2 生成 MET 预测时一致!!!)
WINDOW_SIZE_SEC = 6 # 任务二使用的窗口大小
STEP_SEC = 3        # 任务二使用的窗口步长

# 并行处理进程数
NUM_PROCESSES = mp.cpu_count()

warnings.filterwarnings('ignore', category=FutureWarning)

# --- 绘图函数 (All using English Labels) --- #
def plot_sedentary_summary_distributions(summary_df, save_path=FIGURES_SAVE_PATH):
    """Plot histograms for total duration and bout count per subject."""
    os.makedirs(save_path, exist_ok=True)
    if summary_df.empty:
        print("Info: No summary data available, cannot plot sedentary distributions.")
        return
    plt.figure(figsize=(14, 6))
    # A: Total Duration Histogram
    plt.subplot(1, 2, 1)
    sns.histplot(summary_df['总久坐时长（分钟）'], kde=True)
    plt.title('Distribution of Total Sedentary Duration')
    plt.xlabel('Total Sedentary Duration (minutes)')
    plt.ylabel('Number of Subjects')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    # B: Bout Count Histogram
    plt.subplot(1, 2, 2)
    max_bins = 1
    if not summary_df['久坐次数'].empty:
        max_bins = max(1, int(summary_df['久坐次数'].max()))
    sns.histplot(summary_df['久坐次数'], kde=False, bins=max_bins)
    plt.title('Distribution of Sedentary Bout Count')
    plt.xlabel('Number of Sedentary Bouts')
    plt.ylabel('Number of Subjects')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "sedentary_summary_histograms.png"), dpi=300, bbox_inches='tight') # Renamed file slightly
    plt.close()

def plot_sedentary_bout_duration_distribution(all_bouts_df, save_path=FIGURES_SAVE_PATH):
    """Plot histogram for the duration of individual sedentary bouts."""
    os.makedirs(save_path, exist_ok=True)
    if all_bouts_df.empty:
        print("Info: No sedentary bouts detected, cannot plot duration distribution.")
        return
    plt.figure(figsize=(10, 6))
    sns.histplot(all_bouts_df['久坐持续时长（分钟）'], kde=True)
    plt.title('Distribution of Single Sedentary Bout Duration')
    plt.xlabel('Single Sedentary Bout Duration (minutes)')
    plt.ylabel('Frequency')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "sedentary_bout_duration_histogram.png"), dpi=300, bbox_inches='tight') # Renamed file slightly
    plt.close()

# --- NEW Plotting Functions --- #

def plot_total_sedentary_time_boxplot(summary_df, save_path=FIGURES_SAVE_PATH):
    """Plot boxplot for total sedentary time per subject."""
    os.makedirs(save_path, exist_ok=True)
    if summary_df.empty or '总久坐时长（分钟）' not in summary_df.columns:
        print("Info: No summary data available for total duration boxplot.")
        return
    plt.figure(figsize=(8, 6))
    sns.boxplot(y=summary_df['总久坐时长（分钟）'])
    plt.title('Boxplot of Total Sedentary Duration per Subject')
    plt.ylabel('Total Sedentary Duration (minutes)')
    plt.xlabel('All Subjects')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "total_sedentary_duration_boxplot.png"), dpi=300, bbox_inches='tight')
    plt.close()

def plot_sedentary_bout_count_boxplot(summary_df, save_path=FIGURES_SAVE_PATH):
    """Plot boxplot for the number of sedentary bouts per subject."""
    os.makedirs(save_path, exist_ok=True)
    if summary_df.empty or '久坐次数' not in summary_df.columns:
        print("Info: No summary data available for bout count boxplot.")
        return
    plt.figure(figsize=(8, 6))
    sns.boxplot(y=summary_df['久坐次数'])
    plt.title('Boxplot of Sedentary Bout Count per Subject')
    plt.ylabel('Number of Sedentary Bouts')
    plt.xlabel('All Subjects')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "sedentary_bout_count_boxplot.png"), dpi=300, bbox_inches='tight')
    plt.close()

def plot_sedentary_start_hour_distribution(all_bouts_df, save_path=FIGURES_SAVE_PATH):
    """Plot histogram/density of sedentary bout start hours across all bouts."""
    os.makedirs(save_path, exist_ok=True)
    if all_bouts_df.empty or '久坐开始时间' not in all_bouts_df.columns:
        print("Info: No bout data available for start hour distribution plot.")
        return
    # Extract hour from the datetime column
    all_bouts_df['start_hour'] = pd.to_datetime(all_bouts_df['久坐开始时间']).dt.hour
    plt.figure(figsize=(12, 6))
    sns.histplot(all_bouts_df['start_hour'], bins=24, kde=False)
    plt.title('Distribution of Sedentary Bout Start Times (Hour of Day)')
    plt.xlabel('Hour of Day (0-23)')
    plt.ylabel('Number of Sedentary Bouts Started')
    plt.xticks(range(0, 24))
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "sedentary_start_hour_distribution.png"), dpi=300, bbox_inches='tight')
    plt.close()

def plot_duration_vs_count_scatter(summary_df, save_path=FIGURES_SAVE_PATH):
    """Plot scatter of total sedentary duration vs. bout count per subject."""
    os.makedirs(save_path, exist_ok=True)
    if summary_df.empty or '总久坐时长（分钟）' not in summary_df.columns or '久坐次数' not in summary_df.columns:
        print("Info: Insufficient data for duration vs. count scatter plot.")
        return
    plt.figure(figsize=(10, 6))
    sns.scatterplot(data=summary_df, x='久坐次数', y='总久坐时长（分钟）', alpha=0.7)
    plt.title('Total Sedentary Duration vs. Number of Bouts per Subject')
    plt.xlabel('Number of Sedentary Bouts')
    plt.ylabel('Total Sedentary Duration (minutes)')
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "sedentary_duration_vs_count_scatter.png"), dpi=300, bbox_inches='tight')
    plt.close()


# --- 主处理函数 ---
def process_subject_sedentary_data(args):
    """处理单个受试者数据，识别久坐行为"""
    subject_pid_str = args['subject_id'] # e.g., 'P101'
    subject_tid_str = f"T{subject_pid_str[1:]}" # e.g., 'T101' (用于输出)

    detected_bouts = [] # 存储该受试者检测到的久坐片段

    # 1. 加载该受试者的 MET 预测结果
    met_pred_file = os.path.join(TASK2_PREDICTIONS_DIR, f"{subject_tid_str}_predictions.csv")
    if not os.path.exists(met_pred_file):
        # print(f"警告: 找不到 {subject_tid_str} 的 MET 预测文件，跳过。")
        return detected_bouts # 返回空列表
    try:
        met_preds_df = pd.read_csv(met_pred_file)
        met_preds_df['timestamp'] = pd.to_datetime(met_preds_df['timestamp'])
        met_preds_df = met_preds_df.sort_values('timestamp')
    except Exception as e:
        print(f"错误: 读取 MET 预测文件 {met_pred_file} 失败: {e}")
        return detected_bouts

    if met_preds_df.empty or len(met_preds_df) < 2:
        # print(f"信息: {subject_tid_str} 的 MET 预测数据过少或为空。")
        return detected_bouts

    # 2. 筛选静态且非睡眠窗口 (1.0 <= MET < 1.6)
    static_non_sleep_windows = met_preds_df[
        (met_preds_df['predicted_met'] >= MET_SLEEP_THRESHOLD) &
        (met_preds_df['predicted_met'] < MET_STATIC_THRESHOLD)
    ].copy()

    if static_non_sleep_windows.empty:
        # print(f"信息: {subject_tid_str} 未检测到静态非睡眠窗口 (MET 在 [{MET_SLEEP_THRESHOLD}, {MET_STATIC_THRESHOLD}))。")
        return detected_bouts

    # 3. 识别连续的静态非睡眠窗口
    static_non_sleep_windows['time_diff'] = static_non_sleep_windows['timestamp'].diff().dt.total_seconds()
    # 第一个窗口的 time_diff 是 NaN，我们视为新片段的开始
    # 如果时间差大于步长 STEP_SEC (允许一点点误差)，也视为新片段开始
    # 加上一点误差容忍度，例如 0.5 * STEP_SEC
    static_non_sleep_windows['new_segment'] = (static_non_sleep_windows['time_diff'].isna()) | (static_non_sleep_windows['time_diff'] > STEP_SEC + 0.5 * STEP_SEC)
    static_non_sleep_windows['segment_id'] = static_non_sleep_windows['new_segment'].cumsum()

    # 4. 计算每个连续静态非睡眠片段的时长并筛选久坐行为
    segment_groups = static_non_sleep_windows.groupby('segment_id')

    for segment_id, group in segment_groups:
        if len(group) < 2: # 单个窗口不可能构成久坐
            continue

        start_time = group['timestamp'].iloc[0]
        # 结束时间是最后一个窗口的开始时间 + 窗口大小
        end_time = group['timestamp'].iloc[-1] + timedelta(seconds=WINDOW_SIZE_SEC)
        # 持续时长 = 结束时间 - 开始时间
        duration_seconds = (end_time - start_time).total_seconds()

        # 检查是否超过久坐阈值
        if duration_seconds >= SEDENTARY_DURATION_SECONDS:
            detected_bouts.append({
                '志愿者ID': subject_tid_str,
                '久坐开始时间': start_time,
                '久坐结束时间': end_time,
                '久坐持续时长（分钟）': duration_seconds / 60.0
            })

    return detected_bouts

# --- 主函数 ---
def main():
    start_overall = time.time()
    os.makedirs(RESULTS_DIR, exist_ok=True)
    os.makedirs(FIGURES_SAVE_PATH, exist_ok=True)

    # --- 1. 加载测试集元数据 --- (虽然本次任务不直接用age/sex，但加载以获取受试者列表)
    print("加载测试集元数据...")
    try:
        test_metadata = pd.read_csv(METADATA_TEST)
        if 'pid' not in test_metadata.columns:
            raise ValueError(f"测试元数据 {METADATA_TEST} 缺少 'pid' 列。")
        subject_pids = test_metadata['pid'].unique() # 获取所有测试受试者的 'Pxxx' ID
        print(f"找到 {len(subject_pids)} 个测试受试者。")
    except FileNotFoundError as e:
        print(f"错误: 无法找到测试元数据文件 {e.filename}。请检查路径配置: {DATA_DIR_TEST}")
        sys.exit(1)
    except ValueError as e:
         print(f"错误: 测试元数据文件格式不正确。{e}")
         sys.exit(1)

    # --- 2. 准备并行处理任务 --- #
    print(f"\n开始并行处理 {len(subject_pids)} 个受试者的久坐行为检测 ({NUM_PROCESSES} 个进程)...")
    tasks = [{'subject_id': pid} for pid in subject_pids]

    # --- 3. 并行处理 --- #
    all_detected_bouts_list = []
    with mp.Pool(processes=NUM_PROCESSES) as pool:
        for result_list in tqdm(pool.imap_unordered(process_subject_sedentary_data, tasks), total=len(tasks), desc="检测久坐行为"):
            if result_list: # 如果列表非空
                all_detected_bouts_list.extend(result_list)

    if not all_detected_bouts_list:
        print("\n警告: 未在任何受试者中检测到时长超过 30 分钟的久坐行为。将生成空的 result_4.xlsx。")
        # 创建一个空的 DataFrame 以便保存文件头
        all_bouts_df = pd.DataFrame(columns=['志愿者ID', '久坐开始时间', '久坐结束时间', '久坐持续时长（分钟）'])
    else:
        all_bouts_df = pd.DataFrame(all_detected_bouts_list)
        # 转换时间列格式以便更好地显示
        all_bouts_df['久坐开始时间'] = pd.to_datetime(all_bouts_df['久坐开始时间'])
        all_bouts_df['久坐结束时间'] = pd.to_datetime(all_bouts_df['久坐结束时间'])
        # 按志愿者 ID 和开始时间排序
        all_bouts_df = all_bouts_df.sort_values(by=['志愿者ID', '久坐开始时间'])

    # --- 5. 生成图表和最终结果表 --- #
    print("\n开始生成分析图表和最终汇总表...")

    # 5.1 计算每个受试者的汇总统计 (用于绘图和最终输出)
    if not all_bouts_df.empty:
        summary_stats = all_bouts_df.groupby('志愿者ID').agg(
            总久坐时长_分钟=('久坐持续时长（分钟）', 'sum'),
            久坐次数=('久坐开始时间', 'count')
        ).reset_index()
        # 重命名列以匹配绘图函数预期
        summary_stats.rename(columns={'总久坐时长_分钟': '总久坐时长（分钟）'}, inplace=True)

        # 对于没有检测到久坐的受试者，补充 0 值
        all_subject_tids = [f"T{pid[1:]}" for pid in subject_pids]
        summary_stats_all = pd.DataFrame({'志愿者ID': all_subject_tids})
        summary_stats_all = pd.merge(summary_stats_all, summary_stats, on='志愿者ID', how='left').fillna(0)

        # 确保次数为整数类型
        summary_stats_all['久坐次数'] = summary_stats_all['久坐次数'].astype(int)

        # 绘制图表 A, B (基于汇总)
        plot_sedentary_summary_distributions(summary_stats_all)
        # 绘制图表 C (基于所有检测到的久坐事件)
        plot_sedentary_bout_duration_distribution(all_bouts_df)
        print("图表生成完毕。")
    else:
        # 如果未检测到任何久坐行为，为所有受试者创建汇总记录
        print("信息: 未检测到任何符合条件的久坐行为 (MET 在 [1.0, 1.6) 且持续 > 30 分钟)。")
        all_subject_tids = [f"T{pid[1:]}" for pid in subject_pids]
        summary_stats_all = pd.DataFrame({
            '志愿者ID': all_subject_tids,
            '总久坐时长（分钟）': 0.0,
            '久坐次数': 0
        })
        print("信息: 未检测到久坐行为，无法生成相关图表。")

    # --- Plotting Section --- #
    try:
        # Plots based on summary_stats_all
        plot_sedentary_summary_distributions(summary_stats_all, FIGURES_SAVE_PATH) # Histograms
        plot_total_sedentary_time_boxplot(summary_stats_all, FIGURES_SAVE_PATH) # New Boxplot
        plot_sedentary_bout_count_boxplot(summary_stats_all, FIGURES_SAVE_PATH) # New Boxplot
        plot_duration_vs_count_scatter(summary_stats_all, FIGURES_SAVE_PATH) # New Scatter

        # Plots based on all_bouts_df
        plot_sedentary_bout_duration_distribution(all_bouts_df, FIGURES_SAVE_PATH) # Histogram
        plot_sedentary_start_hour_distribution(all_bouts_df, FIGURES_SAVE_PATH) # New Histogram

        print(f"Plots saved to {FIGURES_SAVE_PATH}")
    except Exception as plot_err:
        print(f"Error during plot generation: {plot_err}")
    # --- End Plotting Section --- #

    # 5.2 添加久坐提醒列
    summary_stats_all['久坐提醒'] = np.where(summary_stats_all['久坐次数'] > 0, 'Sedentary', 'Normal') # Use English labels

    # 5.3 准备最终输出的 DataFrame
    # 选择并排序列
    final_summary_df = summary_stats_all[['志愿者ID', '总久坐时长（分钟）', '久坐次数', '久坐提醒']].copy()
    # 保留4位小数
    final_summary_df['总久坐时长（分钟）'] = final_summary_df['总久坐时长（分钟）'].round(4)
    # 按志愿者ID排序
    final_summary_df = final_summary_df.sort_values(by='志愿者ID')

    # --- 6. 保存最终汇总结果文件 (result_4_summary.xlsx) --- #
    print(f"\n保存每位志愿者的久坐总结到 {OUTPUT_SUMMARY_FILE}...")
    try:
        # 保存汇总表，而非原始事件列表
        final_summary_df.to_excel(OUTPUT_SUMMARY_FILE, index=False, engine='openpyxl')
        print(f"志愿者久坐总结已保存到 {OUTPUT_SUMMARY_FILE}")
    except ImportError:
        print("警告: 未安装 'openpyxl'。请运行 'pip install openpyxl' 来支持 .xlsx 文件写入。尝试保存为 CSV。")
        csv_output_file = OUTPUT_SUMMARY_FILE.replace('.xlsx', '.csv')
        try:
             final_summary_df.to_csv(csv_output_file, index=False)
             print(f"志愿者久坐总结已保存为 CSV 文件: {csv_output_file}")
        except Exception as e_csv:
             print(f"保存为 CSV 时也出错: {e_csv}")
    except Exception as e_excel:
        print(f"保存汇总 Excel 文件时出错: {e_excel}")

    end_overall = time.time()
    print(f"\n--- 任务四处理完成 --- 总耗时: {end_overall - start_overall:.2f} 秒 --- ({len(final_summary_df[final_summary_df['久坐提醒'] == 'Sedentary'])} 人检测到久坐) ---")

# --- 启动进入点 ---
if __name__ == "__main__":
    try:
        print(f"=== 任务四脚本启动时间: {datetime.now()} ===")
        sys.stdout.flush()

        # Windows下多进程支持
        if sys.platform == 'win32':
            mp.freeze_support()

        main()

        print(f"=== 程序正常结束时间: {datetime.now()} ===")
        sys.stdout.flush()
    except Exception as e:
        print("\n=== 程序异常终止 ===")
        print(f"错误信息: {str(e)}")
        import traceback
        print(traceback.format_exc())
        sys.stdout.flush() 